import re

import docx
import typing
from typing import List
import os
import json
import codecs
from docx import Document
import random

color_set = set()
font_set = set()
size_set = set()
bold_set = set()


def read_doc(fname):
    # Document 类，不仅可以新建word文档，也可以打开一个本地文档
    doc = Document(fname)  # 想获取的文档文件名，这里是相对路径。
    # sections = doc.sections  # 获取文档章节信息
    # print(len(sections))

    paragraphs = []
    # print(len(paragraphs))
    chapter_title = None
    for i, p in enumerate(doc.paragraphs):
        texts, title, paragraph_type = process_paragraph(p)
        if len(texts) == 0:
            continue
        if title is not None:
            chapter_title = title
        if paragraph_type == "exegesis":
            paragraphs[-1].add_exegesis(texts)
        elif paragraph_type == "comment" and paragraphs[-1].pairs[0]["text"][0] == chapter_title:
            paragraphs[-1].pairs[0]["comment"].extend([t.text for t in texts])
        else:
            paragraph = Paragraph(texts, chapter_title)
            paragraphs.append(paragraph)
    return paragraphs


class Text:
    def __init__(self, t: str, color: str, size: str, bold: bool):
        self.text = t.strip()
        self.color = color
        self.size = size
        self.bold = bold
        self.role = None
        self.set_role()

    def set_role(self):
        if self.size == "196850" and self.color == "000000":
            self.role = "title"
        elif self.size == "133350" and self.color == "202122":
            if self.text.startswith("（") or self.text.startswith("("):
                self.role = "comment"
            else:
                if not self.text.startswith("疏"):
                    # print(self.text)
                    pass
                self.role = "exegesis"
        elif self.size == "152400" and self.color == "000080":
            self.role = "classic"
        else:
            pass
            # print(self.size, self.color, self.text)
            # raise Exception("wrong size and color")


class Paragraph:
    def __init__(self, texts: List[Text], title):
        self.title = title
        self.pairs = []
        pair = {"text": [], "comment": []}
        self.exegesis = []
        for i in range(len(texts)):
            print(texts[i].text)
            if texts[i].role == "classic" or texts[i].role == "title":
                if len(pair["text"]) != 0:
                    self.pairs.append(pair)
                    pair = {"text": [], "comment": []}
                pair["text"].append(texts[i].text)
            else:
                pair["comment"].append(re.sub(r'[（）]', "", texts[i].text))
        self.pairs.append(pair)
        for pair in self.pairs:
            pair["text"] = "".join(pair["text"])
            pair["comment"] = "".join(pair["comment"])

    def add_exegesis(self, texts):
        for text in texts:
            self.exegesis.append(text.text)

    def to_json(self):
        return {"title": self.title, "pairs": self.pairs, "exegesis": self.exegesis}


"""
    卷标题 ： 000000 196850
    注释文本：202122 133350
    原文：   000080 152400
    疏：    000080 152400
"""


def process_paragraph(p):
    # paragraph对象里还有更小的run对象，run对象才包含了段落对象的文字信息。
    texts = []
    chapter_name = None
    paragraph_type = None
    for r in p.runs:
        font = r.font.name  # 字体均为宋体
        font_set.add(font)
        if font == "Georgia":
            chapter_name = r.text
            paragraph_type = "title"
        size = str(r.font.size)
        size_set.add(size)
        color = str(r.font.color.rgb)
        color_set.add(color)
        bold = r.font.bold
        bold_set.add(bold)
        # print(r.text, color, size, bold)
        t = Text(r.text, color, size, bold)
        if t.role == "classic":
            paragraph_type = "classic"
        if paragraph_type is None:
            if t.role == "exegesis":
                paragraph_type = "exegesis"
            elif t.role == "comment":
                paragraph_type = "comment"
        if t.role == "title":
            paragraph_type = "title"
        texts.append(t)
    return texts, chapter_name, paragraph_type


def build_train_data_for_bge(path="mengzi.jsonl", negtive_sample_rate=2):
    """
    这里不应该用句对句数据训练，因为会泄露，但是可以用段对段数据训练
    :param path:
    :param negtive_sample_rate:
    :return:
    """

    def merge_paragraph(paragraph: dict, info_type: str):
        if info_type == "classic":
            return "".join([p["text"] for p in paragraph["pairs"]])
        elif info_type == "comment":
            return "".join([p["comment"] for p in paragraph["pairs"]])
        elif info_type == "exegesis":
            return "".join(paragraph["exegesis"])

    def write_line(s1, s2, s1_type, s2_type, neg_paragraphs, writer):
        line_data = {"query": s1, "pos": [s2],
                     "neg": [merge_paragraph(json.loads(p), s2_type) for p in neg_paragraphs]}
        writer.write(json.dumps(line_data, ensure_ascii=False) + "\n")
        line_data = {"query": s2, "pos": [s1],
                     "neg": [merge_paragraph(json.loads(p), s1_type) for p in neg_paragraphs]}
        writer.write(json.dumps(line_data, ensure_ascii=False) + "\n")

    lines = open(path).readlines()
    with open("mengzi_exegesis_bge_train.jsonl", "w") as writer:
        for line in lines:
            paragraph = json.loads(line)
            s1 = merge_paragraph(paragraph, "classic")
            s2 = merge_paragraph(paragraph, "comment")
            s3 = merge_paragraph(paragraph, "exegesis")
            neg_paragraphs = random.sample(lines, negtive_sample_rate)
            if len(s1) != 0 and len(s2) != 0:
                write_line(s1, s2, "classic", "comment", neg_paragraphs, writer)
            if len(s1) != 0 and len(s3) != 0:
                write_line(s1, s3, "classic", "exegesis", neg_paragraphs, writer)
            if len(s2) != 0 and len(s3) != 0:
                write_line(s2, s3, "comment", "exegesis", neg_paragraphs, writer)


def write_json(paragraphs: List[Paragraph]):
    with codecs.open("xiaojing.jsonl", "w", "utf-8") as writer:
        for paragraph in paragraphs:
            # print(paragraph.to_json())
            writer.write(json.dumps(paragraph.to_json(), ensure_ascii=False) + "\n")


# 按间距中的绿色按钮以运行脚本。
if __name__ == '__main__':
    build_train_data_for_bge()
